% Reconstruction_tiled_imput.m
%
% author: Kevin O'Holleran - ko311@cam.ac.uk
%
% Description: 2D Structured Illumination reconstruction code. This script
% requires tiled input image (phase of sinusoidal excitation along colums and
% sinusoid orientation along rows). This code was written and tested
% primarly using 512x512 square images - larger images will still work but
% the fringe registration may take a long time. Use 512x512 to measure
% fringe parameters to save time. The reconstruction largely follows the
% method described in Gustafsson 2008 "Three-Dimensional Resolution Doubling in
% Wide-Field Fluorescence Microscopy by Structured Illumination"


%% Input file details
dataPath = './Data/';
filePrefix = 'ps_speck_green8_p12_tiled';
fileSuffix = '.tif';
%
% Physical parameters of the system
% 
dx = 0.04; % pixel size of the image [um] (field of view [um]/number of pixels in image)
n_immersion = 1.515; % refractive index of immersion oil
lambda = 0.520; % centre wavelength for fluorescence emission
NA = 1.3; % numerical aperture of lens
%
% Image width (assumes square image)
width = 512;
% number of phase steps and orientations
nPhases = 3; %no. phase steps
nOrientations = 3; % no. orientations
% x and y coord matrices for performing later calculations
xs = repmat(1:(width),(width),1);
ys = repmat(transpose(1:(width)),1,(width));
% x, y coord arrays for super-resolved image (double resolution)
x = repmat(1:2*width,(2*width),1);
y = repmat(transpose(1:2*width),1,2*width);

% Boolean that switches whether or not to find fringe params or use
% constants loaded from file. The calculated fringe params will be stored
% after calculation in the same directory as the image data.
find_fringe_params = true;
fringeParamsFilename = './SIM_fringe_params';
fringeParamsFilePath = [dataPath fringeParamsFilename];

% default fringe params col 1 = spatial freq, col 2 = orientaion (radians)
if find_fringe_params
    fringeParams = zeros(nOrientations,2);
else
    load(fringeParamsFilePath);
end

%% Determining MTF

pupil_file = 'pupil_data.mat';
measured_MTF_flag = false; % set it to true if measured data is used. 

if measured_MTF_flag % measured MTF
    load([dataPath pupil_file]);
    otf0 =  OTF2D_measured(PUPi,dx,lambda,NA,2*width);
    otf0_small = otf0(256:767,256:767);
else % ideal MTF
    % Set incoherent cutoff (as a ratio of half array size) - this is used as this
    % reconstruction file simply calculates the ideal incoherent OTF with a
    % circular pupil. One way to estimate the cutoff is to perform a FFT of a
    % widefield image of fluorescent beads and measure the radius (in
    % pixels) from the origin (k_x, k_y =0) to the visible passband extent.
    % r0 = radius in pixels/width
    r0 = 150/width;
    % Create zero order OTF matrix at reconstruction resolution
    n = 2*width;
    otf0 = OTF2D(r0,n);
    otf0_small = OTF2D(r0,width);
end

%% Read raw data and reshape
raw_data = imread([dataPath filePrefix fileSuffix],'tif');

% Turn tiled data into a 3D array and apply tukey window to each plane to
% avoid issues with data at edge of images and the FFTs in the
% reconstruction.
iStack = zeros(width,width,nPhases*nOrientations);
for i = 0:(nPhases*nOrientations-1)
    yi = width*mod(i,nPhases)+1;
    xi = width*floor(i/nPhases)+1;
    temp = raw_data(xi:(xi+width-1),yi:(yi+width-1));
    temp = twWindow2D(double(temp));
    iStack(:,:,i+1) = temp;
end

%% Construct widefield image
% This is generated from raw data by summing them and padding in
% Fourier space to generate a widefield image at same pixel size as
% reconstruction.
I_widefield_high_res = sum(iStack,3);
I_widefield_high_res = fft2(I_widefield_high_res);
init_temp = zeros(2*width,2*width);
init_temp((width/2+1):(3*width/2),(width/2+1):(3*width/2)) = fftshift(I_widefield_high_res);
I_widefield_high_res = abs(ifft2(fftshift(init_temp)));

%% Reconstruction starts here

% initialise zero, positive and negative first orders - each orientation's
% data will be stored along the 3rd dimension
k0 = zeros(2*width,2*width,nOrientations);
kp1 = zeros(2*width,2*width,nOrientations);
kn1 = zeros(2*width,2*width,nOrientations);
orders = 2;
otfs = zeros(2*width,2*width,nOrientations);
% initialise complex contrast values for each orientation
ck = zeros(nOrientations,1);

%% Loop over each orientation
% extracting, padding and shifting differnt k-space passbands
for i = 1:nOrientations
    
    % For each orientation solve system of linear equations resulting from
    % three phase images aquired using a single orientation of fringes. If
    % phases are equally spaced over 2*pi then this is the same as performing a
    % Fourier Transform across the phase modulated dimension.
    fft_temp = fft(iStack(:,:,((i-1)*nPhases+1):((i*nPhases))),[],3);
    k0_temp = fft2(fft_temp(:,:,1));
    kp1_temp = fft2(fft_temp(:,:,2));
    
    % Calculate fringe spatial frequency and orientation or index from
    % loaded fringe params mat file.
    if find_fringe_params
        [r_est a_est] = findFringeParams(k0_temp, kp1_temp, 150, otf0_small);
        fringeParams(i,1) = r_est;
        fringeParams(i,2) = a_est;
    else
        r_est = fringeParams(i,1);
        a_est = fringeParams(i,2);
    end
    
    % Pad k-space to accomodate k-space shifts and double resolution in
    % real space
    init_temp = zeros(2*width,2*width);
    init_temp((width/2+1):(3*width/2),(width/2+1):(3*width/2)) = fftshift(k0_temp);
    k0(:,:,i) = fftshift(init_temp);
    init_temp = zeros(2*width,2*width);
    init_temp((width/2+1):(3*width/2),(width/2+1):(3*width/2)) = fftshift(kp1_temp);
    kp1(:,:,i) = fftshift(init_temp);
    
    %k-space translations
    real_space_temp = ifft2(kp1(:,:,i));
    shift_cosine = exp(-1i*2*pi*r_est*(cos(a_est).*x+sin(a_est).*y)/(2*width));
    kp1(:,:,i) = fft2(real_space_temp.*shift_cosine);
    
    % construct negative order from positive first order by rotating 180
    % degrees and taking complex conjugate
    kn1(:,:,i) = conj(rot90(kp1(:,:,i),2));
    kn1(:,:,i) = circshift(kn1(:,:,i),[1 1]);
    
    % Copy and shift OTFs to make components for SIM OTF
    dy = round(r_est*cos(a_est));
    dx = round(r_est*sin(a_est));
    otfp_temp = circshift(otf0,[-dx -dy]);
    otfn_temp = conj(rot90(otfp_temp,2));
    otfn_temp = circshift(otfn_temp,[1 1]);
    otfs(:,:,i) = otfp_temp+otfn_temp;
    % Calculate complex offset of kp1 from k0. This is best explained by
    % looking at Gustafsson 2008 "Three-Dimensional Resolution Doubling in
    % wide in Wide-Field Fluorescence Microscopy by Structured Illumination
    % section: parameter fitting - very bottom of page 4963.
    a_temp = fftshift(k0(:,:,1)).*otfp_temp; % multiply k0_1 by correct translated OTF
    b_temp = fftshift(kp1(:,:,i)).*otf0; % multiply kp1_i by zero order OTF
    % Complex regression to find scalar offsets
    
    dum1_temp = conj(a_temp).*b_temp;
    dum2_temp = abs(a_temp).^2;
    ck(i,1) = 0.5/(sum(dum1_temp(:))/sum(dum2_temp(:)));
end
clear *temp % clears temp arrays

% Save fringeParams matrix in a mat file
save('SIM_fringe_params', 'fringeParams')

% Add together separate components with complex multipliers - these
% constants compensate for the contrast < 1 and any phase shifts relative
% to the zero order. This has implications for noise as noise in the
% super-resolution passbands will be magnified by ck(i)
k_sim = k0(:,:,1);
for i = 1:(nOrientations)
    k_sim  = k_sim + ck(i,1).*kp1(:,:,i)+conj(ck(i,1)).*kn1(:,:,i);
    k_sim  = k_sim + exp(1i*angle(ck(i,1))).*kp1(:,:,i)+exp(-1i*angle(ck(i,1))).*kn1(:,:,i);
end

% Final OTF - the 0.5 is due to 100% contrast corresponding to sidebands
% with a magnitude that is half that of the zero order (positive and
% negative orders sum to give as much 'signal' as the zero order)
otf = otf0+0.5*sum(otfs,3);
% Mask data outside SIM OTF - not necessary but removes the size translated
% squares from the individual k-space shifts.
masked = otf;
masked(masked>0) = 1;
masked = fftshift(masked);
k_sim = masked.*k_sim;

% Again see Gustafsson 2008 3D SIM paper for a full description of this
% filtering process.
% Apodisation function - Cone apodisation out to 2*incoherent cutoff
apodMask = fftshift(triangle(width*2,r0*2*width,0,0));
% Weiner filter
% Weiner filter noise parameter (higer values damped higher spatial
% frequencies
weinerFilterParameter = 0.5;
k_sim_filtered = k_sim./(abs(fftshift(otf)).^2+weinerFilterParameter^2);
k_sim_filtered(isinf(k_sim_filtered))= 0;
% Apodisation
k_sim_filtered = apodMask.*k_sim_filtered;


%% Basic figure showing key images
figure(2)
h = subplot(2,2,1);
imagesc(I_widefield_high_res)
subplot(2,2,2);
sim_image = abs(ifft2(k_sim));
sim_image_filtered = abs(ifft2(k_sim_filtered));
imagesc(sim_image_filtered);
subplot(2,2,3);
k_widefield = fftshift(abs(fft2(I_widefield_high_res)));
imagesc(log(k_widefield));
subplot(2,2,4);
ksr = abs(k_sim_filtered);
ksr = fftshift(ksr);
imagesc(log(ksr))
colormap(jet)

%% Save 16 bit images
n = 2^16-1;
imwrite(uint16(I_widefield_high_res),gray,[dataPath filePrefix 'matlab_widefield.tif'],'tif');
imwrite(uint16(n*sim_image./max(sim_image(:))),gray,[dataPath filePrefix 'matlab_recon.tif'],'tif');
imwrite(uint16(n*sim_image_filtered./max(sim_image_filtered(:))),gray,[dataPath filePrefix 'matlab_recon_filtered.tif'],'tif');

